import unittest

from algo.ratelimiter import TokenBucket


class TestTokenBucket(unittest.TestCase):

    def test_consume(self):
        from time import sleep
        from datetime import datetime

        def now_time():
            return datetime.utcnow().strftime('%Y-%m-%d %H:%M:%S.%f')[:-3]

        bucket = TokenBucket(1, 1)
        is_success = bucket.consume(1)
        print(now_time(), is_success, bucket.tokens)
        self.assertTrue(is_success)
        sleep(0.2)

        for i in range(4):
            is_success = bucket.consume(1)
            print(now_time(), is_success, bucket.tokens)
            self.assertFalse(is_success)
            sleep(0.2)

        is_success = bucket.consume(1)
        print(now_time(), is_success, bucket.tokens)
        self.assertTrue(is_success)
